(*#load "str.cma";;
#load "unix.cma";;
#load "spitree.cmo";;
#load "spiparser.cmo";;
#load "spilexer.cmo";;
#load "ssolver.cma";;
#load "aconv.cmo";;
#load "auxiliary.cmo";;
#load "datquery.cmo";;*)

open Spitree
open Auxiliary
open Datquery
open Printf
open Unix;;

let arguments =
object
  val mutable heap = "ssolver/heap"
  val mutable version = false
  val mutable help = false
  method give_version = version <- true
  method give_help = help <- true
  method print_vers = version
  method print_help = help
  method set_heap h = heap <- h
  method get_heap = heap
end;;

(* Judgment of well-formedness of an environment
 * Input: list env_element env - Environment to judge
 * Output: bool, indicating well-formedness
 *)
exception EnvFailure of string;;
exception TypeDomainInclusionFailure of typ*(env_element list);;
exception DuplicateDeclaration of string*(env_element list);;
exception EffectDomainInclusionFailure of datalog*(env_element list);;
let rec env_judgment env =
  match env with
    | [] -> true
    | Type(x, t)::env_rest ->
(*	Printf.printf "Environment: [%s]\n" (string_of_env env);
	Printf.printf "Fn of %s in domain of [%s]: %B\n" (string_of_typ t) (string_of_env env_rest) (SSet.subset (fn_typ t) (dom env_rest));
	Printf.printf "Name %s not in domain of [%s]: %B\n" x (string_of_env env_rest) (not (SSet.mem x (dom env_rest)));
	Printf.printf "%B\n" (env_judgment env_rest);*)
(*	let good =*)
	let typdominclusion = (SSet.subset (fn_typ t) (dom env_rest)) in
	let notduplicate = (not (SSet.mem x (dom env_rest))) in
	  if not typdominclusion then
(*	    raise (TypeDomainInclusionFailure (sprintf "Free names of type %s not included in environment %s"
						 (string_of_typ t)(string_of_env env_rest)));*)
	    raise (TypeDomainInclusionFailure (t, env_rest));
	  if not notduplicate then
	    raise (DuplicateDeclaration (x, env_rest));
	  (env_judgment env_rest)
(*	if good then true else
	  raise (EnvFailure ("Type declaration " ^ (string_of_env_element (Type(x, t))) ^ " inconsistent with environment " ^ (string_of_env env_rest)))*)
    | Effect(d)::env_rest ->
	let effdominclusion = (SSet.subset (fn_data d) (dom env_rest)) in
	  if not effdominclusion then
	    raise (EffectDomainInclusionFailure (d, env_rest));
	  (env_judgment env_rest);;
(*	if good then true else
	  raise (EnvFailure ("Effect " ^ (string_of_env_element (Effect(d))) ^ "inconsistent with environment " ^ (string_of_env env_rest)));;*)

(* Judgment of the goodness of a message
 * Input:
 * list env_element env - Environment
 * msg msg - Message
 * typ typ - Type
 * Output: bool, indicating if msg is good with type typ in environment env
 *)
exception MsgTypeFailure of string;;
exception MessageTypeMismatch of msg*typ;;
exception NameDomainInclusionFailure of string*(env_element list);;
let rec msg_judgment env msg typ =
  let pair_judgment env m t n u =
    (msg_judgment env m t) &
    (msg_judgment env n u) in
  let encrypt_judgment env m t n ntype =
    (msg_judgment env m t) &
    (msg_judgment env n ntype) in
    try
      match msg with
	| MPair(m, n) ->
	    (match typ with
	       | TPair(x, t, u) ->
		   pair_judgment env m t n (subst_name u x m)
	       | Un ->
		   pair_judgment env m Un n Un
	       | _ -> raise (MsgTypeFailure (sprintf "Expected type Pair or Un, not type %s"
					       (string_of_typ typ))))
	      (*	  | _ -> raise (MessageTypeMismatch (msg, typ)))*)
	| Encr(m, Name(x)) ->
	    (if (typ = Un) then
	       let n, keytype = (Name(x), (nametype x env)) in
		 (match keytype with
		    | Key(t) ->
			encrypt_judgment env m t n keytype
		    | Un ->
			encrypt_judgment env m Un n Un
		    | _ -> raise (MsgTypeFailure (sprintf "Expected type Key or Un for message %s, not type %s"
						    (string_of_msg n) (string_of_typ keytype))))
	     else
	       raise (MsgTypeFailure (sprintf "Expected type Un, not type %s"
					(string_of_typ typ))))
	      (*	   raise (MessageTypeMismatch (msg, typ)))*)
	| Encr(m, n) ->
	    (if (typ = Un) then
	       encrypt_judgment env m Un n Un
	     else
	       raise (MsgTypeFailure (sprintf "Expected type Un, not type %s"
					(string_of_typ typ))))
	      (*	   raise (MessageTypeMismatch (msg, typ)))*)
	| Name(x) ->
	    (*Printf.printf "Name judgment: %s\n" x;*)
	    let ntype = nametype x env in
	    let domaininc = (SSet.mem x (dom env)) in
	    let equality = (type_equal ntype typ) in
	      if not domaininc then
		raise (MsgTypeFailure (sprintf "Name %s not declared in environment %s"
					 x (string_of_env env)));
	      (*	    raise (NameDomainInclusionFailure (x, env));*)
	      if not equality then
		raise (MsgTypeFailure (sprintf "Given type %s does not match type %s from environment."
					 (string_of_typ typ) (string_of_typ ntype)));
	      (env_judgment env)
	| Ok ->
	    (*Printf.printf "Ok judgment\n";*)
	    (match typ with
	       | Un ->
		   env_judgment env
	       | TOk(d) ->
		   let domsubset = SSet.subset (fn_typ typ) (dom env) in
		     if not domsubset then
		       raise (MsgTypeFailure(sprintf "Free names of type %s not a subset of domain of environment %s"
				(string_of_typ typ) (string_of_env env)))
		     else
		   (*	       Printf.printf "Effects of %s in\n" (string_of_reso s);
			       Printf.printf "%s\n" (string_of_env env);*)
		       (querydat (clauses env) (clauses_to_literals d) arguments#get_heap) &
			 (env_judgment env)
	       | _ -> raise (MsgTypeFailure "ok token must have type Un or Ok(S)"))
    with
      | TypeDomainInclusionFailure(t, subenv) ->
	  raise (MsgTypeFailure (sprintf "Message %s:\nAll free names in type %s not included in sub-environment\n%s of environment\n%s"
				   (string_of_msg msg) (string_of_typ t) (string_of_env subenv) (string_of_env env)))
      | DuplicateDeclaration(x, subenv) ->
	  raise (MsgTypeFailure (sprintf "Message %s:\nDuplicate declaration of name \"%s\" in environment\n%s"
				   (string_of_msg msg) x (string_of_env env)))
      | MsgTypeFailure(str) ->
	  raise (MsgTypeFailure (sprintf "Message %s:\n%s" (string_of_msg msg) str))
      | QueryFail(datprog,queryclause) ->
	  raise (MsgTypeFailure (sprintf "Message %s:\nDatalog program %s does not entail %s"
				   (string_of_msg msg) (string_of_datalog datprog) (string_of_hclause queryclause)))
;;

(* Judgment of the goodness of a process
 * Input:
 * list env_element env - Environment
 * proc proc - Process
 * Output: bool, indicating if process proc is good in environment env
 *)
exception ProcTypeFailure of string;;
exception TypeFailure of string;;
let rec proc_judgment env proc =
  let out_judgment env m mtype n t =
    (msg_judgment env m mtype) &
      (msg_judgment env n t) in
  let in_judgment env m mtype x t p =
    (msg_judgment env m mtype) &
      (proc_judgment (Type(x, t)::env) p) in
  let decrypt_judgment env m n ntype y t p =
    (msg_judgment env m Un) &
      (msg_judgment env n ntype) &
      (proc_judgment (Type(y, t)::env) p) in
  let split_judgment env m mtype x t y u p =
    (msg_judgment env m mtype) &
      (proc_judgment (Type(y, u)::Type(x, t)::env) p) in
  let match_judgment env m mtype n t y substu p =
    (msg_judgment env m mtype) &
      (msg_judgment env n t) &
      (proc_judgment (Type(y, substu)::env) p) in
    try
(*      print_endline ((string_of_proc proc) ^ "\n\n");*)
      match proc with
	| Out(Name(x), n) ->
	    let m, chantype = (Name(x), (nametype x env)) in
	      (match chantype with
		 | Ch(t) ->
		     out_judgment env m chantype n t
		 | Un ->
		     out_judgment env m Un n Un
		 | _ -> raise (MsgTypeFailure (sprintf "Expected type Ch or Un for message %s, not type %s"
						  (string_of_msg m) (string_of_typ chantype))))
		(*	   | _ -> raise (ProcTypeFailure ("Channel " ^ (string_of_msg m) ^ "can not have type" ^ (string_of_typ chantype))))*)
	| Out(m, n) ->
	    out_judgment env m Un n Un
	| (In(Name(y), x, p) | InRepl(Name(y), x, p)) ->
	    let m, chantype = (Name(y), (nametype y env)) in
	      (match chantype with
		 | Ch(t) ->
		     in_judgment env m chantype x t p
		 | Un ->
		     in_judgment env m Un x Un p
		 | _ -> raise (MsgTypeFailure (sprintf "Expected type Ch or Un for message %s, not type %s"
						  (string_of_msg m) (string_of_typ chantype))))
		(*	   | _ -> raise (ProcTypeFailure ("Channel " ^ (string_of_msg m) ^ "can not have type" ^ (string_of_typ chantype))))*)
	| (In(m, x, p) | InRepl(m, x, p)) ->
	    in_judgment env m Un x Un p
	| Nu(x, t, p) ->
	    (generative t) &
	      (proc_judgment (Type(x, t)::env) p)
	| PPara(p, q) ->
	    (proc_judgment (List.append (env_of q) env) p) &&
	      (proc_judgment (List.append (env_of p) env) q)
	| PNil -> true
	    (*    | Begin(l, m) ->
		  (SSet.subset (fn_msg m) (dom env)) &
		  (env_judgment env)
		  | End(l, m) ->
		  (SSet.subset (fn_msg m) (dom env)) &
	    (*	(effectinclusion [Reso(l, m)] (effects env)) &*)
		  (env_judgment env)*)
	| Datprog(d) ->
	    let domsubset = SSet.subset (fn_data d) (dom env) in
	      if not domsubset then
		raise (MsgTypeFailure(sprintf "Free names of program %s not a subset of domain of environment %s"
					(string_of_datalog d) (string_of_env env)))
	      else
		(env_judgment env)
	| Expect(l) ->
	    let domsubset = SSet.subset (fn_literal l) (dom env) in
	      if not domsubset then
		raise (MsgTypeFailure(sprintf "Free names of literal %s not a subset of domain of environment %s"
					(string_of_literal l) (string_of_env env)))
	      else
		(querydat (clauses env) [l] arguments#get_heap) &&
		  (env_judgment env)
	| Decrypt(m, y, t1, Name(x), p) ->
	    let n, keytype = (Name(x), (nametype x env)) in
	      (match keytype with
		 | Key(t2) ->
		     if type_equal t1 t2 then
		       decrypt_judgment env m n keytype y t1 p
		     else
		       raise (MsgTypeFailure (sprintf "Decrypt failure: Key type %s incompatible with message type %s"
						 (string_of_typ keytype) (string_of_typ t1)))
		 | Un ->
		     decrypt_judgment env m n Un y Un p
		 | _ -> raise (MsgTypeFailure (sprintf "Expected type Key or Un for message %s, not type %s"
						  (string_of_msg n) (string_of_typ keytype))))
		(*	   | _ -> raise (ProcTypeFailure ("Key " ^ (string_of_msg n) ^ "can not have type" ^ (string_of_typ keytype))))*)
	| Decrypt(m, y, t, n, p) ->
	    decrypt_judgment env m n Un y Un p
	| Split(Name(z), x, t, y, u, p) ->
	    let m, pairtype = (Name(z), (nametype z env)) in
	      (match pairtype with
		 | TPair(x_dep, t_dep, u_dep) ->
		     let subst_u = subst_name u_dep x_dep (Name(x)) in
		     let subst_pairtype = TPair(x, t, subst_u) in
		       if (type_equal t_dep t) then
			 if (type_equal subst_u u) then
			   (split_judgment env m subst_pairtype x t y u p)
			 else
			   raise (MsgTypeFailure (sprintf "Split failure: Type %s incompatible with type %s"
						    (string_of_typ u_dep) (string_of_typ u)))
		       else
			 raise (MsgTypeFailure (sprintf "Split failure: Type %s incompatible with type %s"
						  (string_of_typ t_dep) (string_of_typ t)))
		 | Un ->
		     split_judgment env m Un x Un y Un p
		 | _ -> raise (MsgTypeFailure (sprintf "Expected type Pair or Un for message %s, not type %s"
						  (string_of_msg m) (string_of_typ pairtype))))
		(*	     | _ -> raise (ProcTypeFailure ("Pair " ^ (string_of_msg m) ^ "can not have type " ^ (string_of_typ pairtype))))*)
	| Split(MPair(m1, m2), x, t, y, u, p) ->
	    let m, pairtype = (MPair(m1, m2), TPair(x, t, u)) in
	      split_judgment env m pairtype x t y u p
	| Split(m, x, t, y, u, p) ->
	    split_judgment env m Un x Un y Un p
	| Match(m, mtype, n, y, usub, p) ->
	    (match mtype with
	       | Un ->
		   match_judgment env m Un n Un y Un p
	       | TPair(x, t, u) ->
		   if type_equal (subst_name u x n) usub then
		     (match_judgment env m mtype n t y usub p)
		   else
		     raise (MsgTypeFailure (sprintf "type\n%s does not match substituted type\n%s."
					      (string_of_typ usub) (string_of_typ (subst_name u x n))));
	       | _ ->
		   raise (MsgTypeFailure (sprintf "Expected type Pair or Un for message %s, not type %s"
					     (string_of_msg m) (string_of_typ mtype))))
	      (*	       raise (ProcTypeFailure ("Message " ^ (string_of_msg m) ^ " can not have type " ^ (string_of_typ mtype))))*)
    with
      (*| ProcTypeFailure(str) -> raise (TypeFailure (sprintf "In process:\n %s\n%s" (string_of_proc proc) str))*)
      | MsgTypeFailure(str) -> raise (TypeFailure (sprintf "In process:\n %s\n%s" (string_of_proc proc) str))
      | QueryFail(datprog,queryclause) ->
	  raise (TypeFailure (sprintf "In process %s:\nDatalog program %s does not entail %s"
				(string_of_proc proc) (string_of_datalog datprog) (string_of_hclause queryclause)))
;;



(*effects_in_type (TPair("tnavn", Key(TOk(RPair(Reso("send", Name("a")),Reso("send2", Name("b"))))), TOk(Reso("labl", Name("Moha")))));;

dom [Type("muaah", Un); Type("hehe", TOk(RPair(Reso("heigen",Name("Hmm")),Reso("Gnof", Name("Hnah"))))); Effect(Reso("he", Name("hmm")))];;

msg_judgment [Effect(Reso("send", Name("a"))); Type("noegle", (Key(TOk(Reso("send", Name("a")))))); Type("a", Un)] (Encr(Ok, Name("noegle"))) Un;;
*)

(*let infile = "test.spi";;
let instring = "split pair(hej1,hej2) as (hov:T, he:T); nil"
let msgstring = "({ok}noegle, enlabel({ok}noegle))";;
let procstring = "(nil|(nil|(nil|nil)))";;*)

(*let heap = "ssolver/heap" in*)
let name = "Spi Policy Checker" in
let version = "0.9" in

let help_msg =
  "Usage: " ^ Sys.argv.(0) ^ " [options] filename\n\n" ^
    "Options:\n" ^
    " --solverheap <heapfile>    compiled heap of the Succinct Solver\n" ^
    "                            (default: ssolver/heap)\n" ^
    " --help                     display this help and exit\n" ^
    " --version                  output version information and exit"
in

  if (Array.length Sys.argv) = 1 then
    print_endline (help_msg)
  else
    let infile = Array.fold_left (fun prev_opt this_opt ->
			       match prev_opt with
				 | "" -> this_opt
				 | "--solverheap" ->
				     arguments#set_heap this_opt;
				     ""
				 | "--version" ->
				     arguments#give_version;
				     this_opt
				 | _ ->
				     arguments#give_help;
				     this_opt
			    ) "" (Array.sub Sys.argv 1 ((Array.length Sys.argv) - 1))
    in
      if infile = "--version" then
	arguments#give_version
      else if infile = "--help" then
	arguments#give_help;
      if ((not arguments#print_help) && (not arguments#print_vers)) then
	let filechan = open_in infile in
	let lexbuf = Lexing.from_channel filechan in
	let ast = Spiparser.process Spilexer.spi_tokens lexbuf in
	let (ast_substituted, accum, substMap) = Aconv.proc_subst ast 0 Aconv.SMap.empty in
	let freevars =
	  Aconv.SMap.fold (fun _ name list ->
			     Type(name, Un) :: list) substMap [] in
	  print_endline "Alpha converted process:";
	  print_endline ((string_of_proc ast_substituted) ^ "\n");
	  print_endline "Free names:";
	  Aconv.SMap.iter (Printf.printf "%s: %s\n") substMap;
	  try
	    let judgment = proc_judgment freevars ast_substituted in
	      if judgment then
		Printf.printf "\nCongratulations! This process is robustly safe.\n";
	  with
	    | TypeFailure(str) -> print_endline ("Failure:\n" ^ str)
	    (*| ProcTypeFailure(str) -> print_endline ("Failure:\n" ^ str)
	    | MsgTypeFailure(str) -> print_endline ("Failure:\n" ^ str)*)

      else
	if arguments#print_vers then
	  print_endline (name ^ " " ^ version)
	else
	  print_endline (help_msg);;

(*let infile = Sys.argv.(1);;*)

(*let infile = "signdecrypt.spi";;*)

(* We check for robust safety by initializing all
   free variables to be of type Un *)
